#CNN建立策略和价值网络
#输入为四个平面：1黑棋已下的平面；2白棋已下的平面；3最新的一步棋；4当前执棋者
#输出15*15+1个值

import numpy as np
import tensorflow as tf 

class PolicyValueNet(object):
    #搭建计算图(神经网络)
    def __init__(self,width,height,model_file=None):
        self.board_width=width
        self.board_height=height

        #1输入
        self.input_states=tf.placeholder(tf.float32,shape=[None,4,height,width])
        self.input_state=tf.transpose(self.input_states,[0,2,3,1])
        #输入变为高*宽*平面数

        #2通用卷积层
        #2.1卷积层1，输入[15,15,4]，输出[15,15,32]
        self.conv1=tf.layers.conv2d(self.input_state,filters=32,
                                    kernel_size=[3,3],padding='same',
                                    activation=tf.nn.relu)
        #2.2卷积层2，输入[15,15,32]，输出[15,15,64]
        self.conv2=tf.layers.conv2d(self.conv1,filters=64,
                                    kernel_size=[3,3],padding='same',
                                    activation=tf.nn.relu)
        #2.3卷积层3，输入[15,15,64]，输出[15,15,128]
        self.conv3=tf.layers.conv2d(self.conv2,filters=128,
                                    kernel_size=[3,3],padding='same',
                                    activation=tf.nn.relu)
        
        #策略（动作）和价值网络分离
        #3.1动作过渡层，输入[15，15，128]，输出[15，15，4]
        self.action_conv=tf.layers.conv2d(self.conv3,filters=4,
                                        kernel_size=[1,1],padding='same',
                                        activation=tf.nn.relu)
        #3.2将过渡层摊成一维
        self.action_flatten=tf.reshape(self.action_conv,
                                        shape=[-1,4*height*width])
        #3.3全连接层，输入15*15*4,输出15*15个棋盘位置的概率p的log
        self.action_fc=tf.layers.dense(self.action_flatten,
                                        units=height*width,
                                        activation=tf.nn.log_softmax)
        #4.1价值过渡层，输入[15，15，128]，输出[15，15，2]
        self.evalution_conv=tf.layers.conv2d(self.conv3,filters=2,
                                        kernel_size=[1,1],padding='same',
                                        activation=tf.nn.relu)
        #4.2将过渡层摊成一维
        self.evalution_flatten=tf.reshape(self.evalution_conv,
                                            shape=[-1,2*height*width])
        #4.3全连接层，输入15*15*2，输出1个在[-1,1]之间的价值v
        self.evalution_fc1=tf.layers.dense(self.evalution_flatten,
                                        units=64,activation=tf.nn.relu)
        self.evalution_fc2=tf.layers.dense(self.evalution_fc1,
                                            units=1,activation=tf.nn.tanh)
        #定义损失函数loss = (z - v)^2 - π^T * log(p) +  c||theta||^2
        #1 label:对每个state下的胜利者z（mcts），训练神经网络的标准
        #  mcts_prob:由mcts得出的每个位置的概率π
        self.labels=tf.placeholder(tf.float32,shape=[None,1])
        self.mcts_probs=tf.placeholder(tf.float32,shape=[None,height*width])

        #2 价值损失函数mse (z-v)^2
        self.value_loss=tf.losses.mean_squared_error(self.labels,self.evalution_fc2)

        #3 策略损失函数 -π^T * log(p)
        self.policy_loss=tf.negative(tf.reduce_mean(tf.reduce_sum(
                                    tf.multiply(self.mcts_probs,self.action_fc),1)))

        #4 L2正则化 c||theta||^2，对权重求平方和
        l2_lamda=1e-4
        vars=tf.trainable_variables()
        l2_regulation=l2_lamda*tf.add_n(
            [tf.nn.l2_loss(v) for v in vars if 'bias' not in v.name.lower()])
        #得到tf.add_n([,,,,])，得到各项平方后求和

        #5 损失函数
        self.loss=self.value_loss+self.policy_loss+l2_regulation

        #确定优化器
        self.learning_rate=tf.placeholder(tf.float32)#***
        self.optimizer=tf.train.AdamOptimizer(learning_rate=self.learning_rate)
        self.global_step=tf.Variable(0,name='global_step',trainable=False)
        self.train_op=self.optimizer.minimize(self.loss,global_step=self.global_step)

        #计算策略action_fc信息熵——内在混乱程度 -p^T * log(p)
        self.entropy=tf.negative(tf.reduce_mean(tf.reduce_sum(
                                tf.exp(self.action_fc)*self.action_fc,1)))
        
        #执行计算图，session
        self.session=tf.Session()
        #初始化变量
        init=tf.global_variables_initializer()
        self.session.run(init)
        #保存文件
        self.saver=tf.train.Saver()
        if model_file is not None:
            tf.reset_default_graph()
            self.restore_model(model_file)

    #对战时用单次NN评估，得到（p，v）
    #训练时使用，利用未更新模型，这个函数中没有优化过程
    def policy_value(self,state_batch):
        #输入：a batch of states(s)
        #输出：a batch of action probabilities(p一维) and states value(v)
        log_act_probs,value=self.session.run(
                            [self.action_fc,self.evalution_fc2],
                            feed_dict={self.input_states:state_batch})
        act_probs=np.exp(log_act_probs)
        return act_probs,value
    
    def policy_value_fun(self,board):
        #输入：board
        #中间调用policy_value
        #输出：a list of action probabilities(p) and states value(v)
        legal_positions = board.availables
        current_state = np.ascontiguousarray(board.current_state().reshape(
                -1, 4, self.board_width, self.board_height))
        act_probs, value = self.policy_value(current_state)
        act_probs = zip(legal_positions, act_probs[0][legal_positions])#（position,prob）
        return act_probs, value        
    

    #训练需要计算损失函数，过程中会计算概率和值，同时有优化过程
    def train_step(self, state_batch, mcts_probs, winner_batch, lr):
        #winner_batch是z，mcts_probs是π
        winner_batch = np.reshape(winner_batch, (-1, 1))
        loss, entropy, _, value_loss, policy_loss = self.session.run(
                [self.loss, self.entropy, self.train_op, self.value_loss, self.policy_loss],
                feed_dict={self.input_states: state_batch,
                           self.mcts_probs: mcts_probs,
                           self.labels: winner_batch,
                           self.learning_rate: lr})
        return loss, entropy, value_loss, policy_loss


    #保存文件
    def save_model(self, model_path):
        self.saver.save(self.session, model_path)

    def restore_model(self, model_path):
        self.saver.restore(self.session, model_path)